# 截图宝 - 功能窗体的基类
# 作者： 魏新崇
# 邮箱： weixinchong [at] 163.com
# 项目主页：https://gitee.com/weixinchong/auto-igrab

import sys, os, cv2, imageio, fitz, numpy as np
from PyQt5.QtWidgets import QApplication, QWidget, QGridLayout, QLabel, QPushButton, QSpinBox, QFileDialog
from PyQt5.QtGui import QFont
from PIL import ImageGrab
from psdialog import PointSelDialog

class Window_Base(QWidget):

    def create_spin(self, v0, vmax = 10000, vmin = 0):
        '''创建数字输入框'''
        box = QSpinBox()
        box.setMaximum(vmax)
        box.setMinimum(vmin)
        box.setValue(v0)
        return box

    def __init__(self, title, sc, parent = None):
        super(Window_Base, self).__init__(parent)
        self.setWindowTitle(title)

        # 获得当前路径
        self.cwd = os.path.split(os.path.realpath(__file__))[0].replace('\\', '/')

        # 定义字体
        self.font_title = QFont()
        self.font_title.setBold(True)

        # 设置网格布局
        self.grid = QGridLayout()
        self.grid.setSpacing(10)
        self.setLayout(self.grid)
        self.adjustSize()

        # 添加截图坐标范围相关的控件
        # 1. 创建控件
        label_sc_range = QLabel('截图坐标范围')
        label_sc_range.setFont(self.font_title)
        label_sc_lt = QLabel('左上角：')
        label_sc_rb = QLabel('右下角：')
        self.spin_sc_left = self.create_spin(sc[0])
        self.spin_sc_top = self.create_spin(sc[1])
        self.spin_sc_right = self.create_spin(sc[2])
        self.spin_sc_bottom = self.create_spin(sc[3])
        btn_sc_range_adjust = QPushButton("设定选区")
        btn_sc_range_adjust.clicked.connect(self.get_area)
        # 2. 控件导入布局
        self.grid.addWidget(label_sc_range, 1, 0, 1, 2)
        self.grid.addWidget(label_sc_lt, 2, 0)
        self.grid.addWidget(label_sc_rb, 3, 0)
        self.grid.addWidget(self.spin_sc_left, 2, 1)
        self.grid.addWidget(self.spin_sc_top, 2, 2)
        self.grid.addWidget(self.spin_sc_right, 3, 1)
        self.grid.addWidget(self.spin_sc_bottom, 3, 2)
        self.grid.addWidget(btn_sc_range_adjust, 2, 3)

    def get_params(self):
        '''读取用户输入的参数值'''
        self.sc_left = self.spin_sc_left.value()
        self.sc_top = self.spin_sc_top.value()
        self.sc_right = self.spin_sc_right.value()
        self.sc_bottom = self.spin_sc_bottom.value()

    def get_position(self, o1, o2):
        '''启动获取坐标点对话框'''
        self.pointsel_dialog = PointSelDialog()
        self.pointsel_dialog.setSinglePointObj(o1, o2)
        self.pointsel_dialog.show()

    def get_area(self):
        '''启动获取截图区域对话框'''
        self.pointsel_dialog = PointSelDialog()
        self.pointsel_dialog.setAreaObj(self.spin_sc_left, self.spin_sc_top, self.spin_sc_right, self.spin_sc_bottom)
        self.pointsel_dialog.show()

    def window_capture(self, fName):
        '''截图'''
        im = ImageGrab.grab(bbox = (self.sc_left, self.sc_top, self.sc_right, self.sc_bottom))
        im.save(fName)

    def window_capture_part(self, fName, ydisp):
        '''增补区域截图'''
        im = ImageGrab.grab(bbox = (self.sc_left, self.sc_bottom - ydisp, self.sc_right, self.sc_bottom))
        im.save(fName)

    def sel_pic_path(self):
        '''设置图片的保存地址'''
        dirName = QFileDialog.getExistingDirectory(self, '请选择截取图片的保存地址', self.dir_pic.text())
        if dirName != '':
            self.dir_pic.setText(dirName)

    def sel_video_path(self):
        '''设置视频的保存地址'''
        fName, fType = QFileDialog.getSaveFileName(self, '请选择动画的保存地址', self.dir_video.text(), '视频文件 (*.mp4);;gif动画 (*.gif)')
        if fName != '':
            self.dir_video.setText(fName)

    def sel_pdf_path(self):
        '''设置PDF文档的保存地址'''
        fName, fType = QFileDialog.getSaveFileName(self, '请选择要输出PDF的保存地址', self.dir_pdf.text(), 'PDF文档 (*.pdf)')
        if fName != '':
            self.dir_pdf.setText(fName)

    def init_video(self, vName):
        '''初始化视频部分'''
        self.file_type = vName.split('.')[1]
        if self.file_type == 'mp4':
            size = (self.sc_right - self.sc_left, self.sc_bottom - self.sc_top)
            fourcc = cv2.VideoWriter_fourcc(*'avc1')
            self.videoWriter = cv2.VideoWriter(vName, fourcc, self.fps, size)
        elif self.file_type == 'gif':
            self.videoWriter = []

    def append_frame(self, fName):
        '''将当前图片增加到帧'''
        if self.file_type == 'mp4':
            frame = cv2.imdecode(np.fromfile(fName, dtype = np.uint8), -1)
            self.videoWriter.write(frame)
        elif self.file_type == 'gif':
            frame = imageio.imread(fName)
            self.videoWriter.append(frame)

    def end_video(self, vName):
        '''结束写入视频'''
        if self.file_type == 'mp4':
            self.videoWriter.release()
        elif self.file_type == 'gif':
            imageio.mimsave(vName, self.videoWriter, fps = self.fps)

    def init_pdf(self):
        '''初始化pdf'''
        self.doc = fitz.open()

    def append_pdf_page(self, fName):
        '''将当前图片增加到pdf'''
        imgdoc = fitz.open(fName)
        pdfbytes = imgdoc.convertToPDF()
        imgpdf = fitz.open('pdf', pdfbytes)
        self.doc.insertPDF(imgpdf)
    
    def end_pdf(self, pName):
        '''结束写入pdf'''
        self.doc.save(pName)